// Copyright (c) 2025 Shenzhen Kaihong Digital Industry Development Co., Ltd.
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

use dora_node_api::{self, DoraNode, Event, EventStream};
use eyre::eyre;
use rerun::external::{clap::{self, Parser}, log::{error, info}, re_log::{self},
};
use rerun_tools::{dora_data::DataKind, tensor_data_to_bytes_with_shape};
use std::path::PathBuf;

#[derive(Debug, Parser)]
#[clap(author, version, about)]
struct Args {
    /// The path to rgb image file
    rgb_path: PathBuf,
    /// The path to depth image file
    depth_path: PathBuf,
}

fn run(
    mut node: DoraNode,
    mut events: EventStream,
    rgb_path: PathBuf,
    depth_path: PathBuf,
) -> eyre::Result<()> {
    let mut has_sent_rgb_images = true;
    let mut has_sent_depth_images = true;
    let rgb_image = rerun::Image::from_file_path(rgb_path).expect("failed to load image");
    let depth_image = rerun::Image::from_file_path(depth_path).expect("failed to load image");
    let rgb_tensor_data = rgb_image.data;
    let depth_tensor_data = depth_image.data;
    loop {
        let event = match events.recv() {
            Some(input) => input,
            None => break,
        };
        match event {
            Event::Stop(_) => {
                info!("dora stop");
                return Ok(());
            },
            Event::Reload { operator_id } => {
                return Err(eyre!("unhandled dora reload event: {:?}", operator_id));
            },
            Event::Input {
                id,
                metadata: _,
                data: _,
            } => match id.as_str() {
                "tick" => {
                    if has_sent_rgb_images {
                        let array = tensor_data_to_bytes_with_shape(rgb_tensor_data.clone());
                        node.send_output_bytes(
                            String::from(DataKind::RgbImage.as_output_id()).into(),
                            Default::default(),
                            array.len(),
                            &array,
                        )?;
                        has_sent_rgb_images = false;
                    } else if !has_sent_rgb_images && has_sent_depth_images {
                        let array = tensor_data_to_bytes_with_shape(depth_tensor_data.clone());
                        node.send_output_bytes(
                            String::from(DataKind::DepthImage.as_output_id()).into(),
                            Default::default(),
                            array.len(),
                            &array,
                        )?;
                        has_sent_depth_images = false;
                    }
                }
                other => {
                    error!("unexpected input: {}", other);
                }
            },
            Event::InputClosed { id } => {
                info!("dora input {} closed", id);
                return Ok(());
            },
            Event::Error(msg) => {
                return Err(eyre!("received dora error event: {}", msg));
            },
            _ => {
                error!("unkown dora event: {:?}", event);
            }
        }
    }
    Ok(())
}

#[tokio::main]
async fn main() -> eyre::Result<()> {
    re_log::setup_logging();
    let args = Args::try_parse()?;
    let (node, events) = DoraNode::init_from_env()?;
    run(node, events, args.rgb_path, args.depth_path)?;
    Ok(())
}
